# -*- coding: utf-8 -*-
# @Time    : 2021/8/9 17:54
# @Author  : huangwei
# @File    : method.py
# @Software: PyCharm
from data_info import store_dict, good_dict, sort_list
import numpy as np
import random


def get_data( filepath, ratio ):
    # 取数据
    store_input, good_input, labels, test_store, test_good, test_labels = [], [], [], [], [], []
    sort_input, test_sort = [], []

    with open(filepath) as f:
        lines = f.readlines()
        for line in lines:
            line = line.strip().split(',')
            k = random.uniform(0, 1)
            if k < ratio:
                test_store.append(store_dict[line[0]])
                test_good.append(good_dict[line[1]])
                test_sort.append(sort_list[store_dict[line[0]]])
                test_labels.append(float(line[2]))
            else:
                store_input.append(store_dict[line[0]])
                good_input.append(good_dict[line[1]])
                sort_input.append(sort_list[store_dict[line[0]]])
                labels.append(float(line[2]))

    return store_input, good_input, sort_input, labels, test_store, test_good, test_sort, test_labels


def get_one_hot( data ):
    if data == 0:
        return [1, 0, 0]
    elif data == 1:
        return [0, 1, 0]
    else:
        return [0, 0, 1]


def load_data( store, good, data_sort, label, batch_size, mode='train' ):
    data_length = len(store)
    index_list = list(range(data_length))

    def data_generator():
        if mode == 'train':
            random.shuffle(index_list)

        store_list, good_list, sort_data, label_list = [], [], [], []

        for idx, i in enumerate(index_list):
            store_list.append(store[i])
            good_list.append(good[i])
            sort_data.append(get_one_hot(data_sort[i]))
            label_list.append(label[i])

            if len(store_list) == batch_size:
                store_arr = np.array(store_list)
                good_arr = np.array(good_list)
                sort_arr = np.array(sort_data)
                label_arr = np.array(label_list)

                yield store_arr, good_arr, sort_arr, label_arr
                store_list, good_list, sort_data, label_list = [], [], [], []

    return data_generator


def data_gen( filepath, test_ratio, batch_size=1024 ):
    store_input, good_input, sort_input, labels, test_store, test_good, test_sort, test_labels = get_data(filepath,
                                                                                                          test_ratio)
    train_loader = load_data(store_input, good_input, sort_input, labels, batch_size)
    test_loader = load_data(test_store, test_good, test_sort, test_labels, batch_size, 'test')

    return train_loader, test_loader
